from wcwidth import wcswidth
import numpy as np
import cv2
from sklearn.cluster import KMeans

def get_luminance(color):
    """计算颜色亮度（ITU-R BT.709标准）"""
    r, g, b = (int(color[i+1:i+3], 16) for i in (0, 2, 4))  # 转换HEX颜色到RGB
    return 0.2126 * r + 0.7152 * g + 0.0722 * b


def get_contrast_colors(bg_color):
    """获取自适应对比色"""
    brightness = get_luminance(bg_color)

    # 动态颜色选择逻辑
    if brightness < 60:
        text = "#FFFFFF"
        stroke = "#000000" if brightness > 30 else "#FFFFFF"
    elif brightness > 200:
        text = "#000000"
        stroke = "#FFFFFF" if brightness < 225 else "#000000"
    else:
        text = "#000000" if brightness < 127 else "#FFFFFF"
        stroke = "#FFFFFF" if brightness < 127 else "#000000"

    return text, stroke


def get_color_from_class(cls):
    # YOLO标准16色色板（与网页2、网页4配色一致）
    yolo_color_palette = [
        (255, 0, 255),   # 0. 品红色
        (0, 0, 255),     # 1. 红色
        (0, 255, 0),     # 2. 绿色
        (255, 0, 0),     # 3. 蓝色
        (0, 255, 255),   # 4. 黄色
        (255, 165, 0),   # 5. 橙色
        (75, 0, 130),    # 6. 靛蓝
        (255, 20, 147),  # 7. 深粉色
        (0, 255, 127),   # 8. 春绿色
        (255, 255, 0),   # 9. 青色
        (128, 0, 128),   # 10. 紫色
        (255, 99, 71),   # 11. 番茄红
        (60, 179, 113),  # 12. 海洋绿
        (238, 130, 238), # 13. 紫罗兰
        (255, 215, 0),   # 14. 金色
        (0, 191, 255)    # 15. 天空蓝
    ]
    return yolo_color_palette[cls % len(yolo_color_palette)]


def chinese_pad(text, width, align='^'):
    """支持中文的字符串对齐函数"""
    current_len = wcswidth(text)
    remaining = width - current_len
    if remaining <= 0:
        return text
    if align == '<':
        return text + ' ' * remaining
    elif align == '>':
        return ' ' * remaining + text
    else:
        left = remaining // 2
        right = remaining - left
        return ' ' * left + text + ' ' * right


def cluster_largest_region(image_array, n_clusters=3):
    """
    使用 HSV + 邻域差值进行聚类，仅保留图像中心点所属类的腐蚀轮廓。

    参数:
        image_array: 输入图像（BGR）
        n_clusters: 聚类类别数

    返回:
        damage_area: int，区域像素数量
        contours: list[np.ndarray]，腐蚀后轮廓
    """
    # erode_iter = 2
    # kernel_size = 3
    h, w = image_array.shape[:2]
    if h == 0 or w == 0:
        return 0, []

    # 1. 转 HSV
    hsv = cv2.cvtColor(image_array, cv2.COLOR_BGR2HSV)
    h_ch, s_ch, v_ch = hsv[..., 0], hsv[..., 1], hsv[..., 2]

    rgb = cv2.cvtColor(image_array, cv2.COLOR_BGR2RGB)
    r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]
    # 2. 计算差分 (与上下左右邻居的差值平均)
    # def mean_neighbor_diff(channel):
    #     padded = np.pad(channel, pad_width=1, mode='reflect')
    #     center = padded[1:-1, 1:-1].astype(np.float32)
    #     neighbors = (
    #         np.abs(padded[0:-2, 1:-1] - center) +
    #         np.abs(padded[2:,   1:-1] - center) +
    #         np.abs(padded[1:-1, 0:-2] - center) +
    #         np.abs(padded[1:-1,  2:] - center)
    #     ) / 4.0
    #     return neighbors
    #
    # dh = mean_neighbor_diff(h_ch)
    # ds = mean_neighbor_diff(s_ch)
    # dv = mean_neighbor_diff(v_ch)

    # 3. 构造特征向量: [h, s, v, Δh, Δs, Δv]

    features = np.stack([
        h_ch.ravel(), s_ch.ravel(), v_ch.ravel(),
        r.ravel(), g.ravel(), b.ravel(),
        # dh.ravel(), ds.ravel(), dv.ravel()
    ], axis=1)

    # 4. 聚类
    kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=20)
    labels = kmeans.fit_predict(features)
    label_map = labels.reshape(h, w)

    # 5. 中心标签
    center_y, center_x = h // 2, w // 2
    # center_label = label_map[center_y, center_x]
    half_window = 5 // 2
    y_start = max(0, center_y - half_window)
    y_end = min(h, center_y + half_window + 1)
    x_start = max(0, center_x - half_window)
    x_end = min(w, center_x + half_window + 1)
    center_region = label_map[y_start:y_end, x_start:x_end]
    unique, counts = np.unique(center_region, return_counts=True)
    center_label = unique[np.argmax(counts)]

    # 6. 生成掩码
    raw_mask = (label_map == center_label).astype(np.uint8) * 255

    # 7. 连通区域分析，剔除面积小于50的区域
    num_labels, labels_mask, stats, _ = cv2.connectedComponentsWithStats(raw_mask, connectivity=8)

    filtered_mask = np.zeros_like(raw_mask, dtype=np.uint8)
    for i in range(1, num_labels):  # 跳过背景（标签0）
        area = stats[i, cv2.CC_STAT_AREA]
        if area >= 50:
            filtered_mask[labels_mask == i] = 255  # 保留该区域

    mask = filtered_mask  # 最终掩码
    # 7. 腐蚀
    # kernel = np.ones((kernel_size, kernel_size), np.uint8)
    # mask_eroded = cv2.erode(mask, kernel, iterations=erode_iter)

    # 8. 轮廓
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    damage_area = int(np.sum(mask > 0))
    return damage_area, contours


def estimate_damage_area_by_clustering(image_array, box):
    """
    通过KMeans聚类估算检测框内的损毁面积，使用(x, y, 灰度)作为聚类特征。
    较小的聚类簇被认为是损毁区域。
    """
    x1, y1, x2, y2 = box
    crop = image_array[y1:y2, x1:x2]  # 提取框内区域
    h, w = crop.shape[:2]
    if h == 0 or w == 0:
        return 0  # 空区域直接返回0

    # 转灰度图（用于亮度聚类）
    gray = np.dot(crop[..., :3], [0.2989, 0.5870, 0.1140])  # shape: (h, w)

    # 构造特征：位置 + 亮度
    yy, xx = np.meshgrid(np.arange(h), np.arange(w), indexing='ij')  # 坐标网格
    features = np.stack([xx.ravel(), yy.ravel(), gray.ravel()], axis=1)  # shape: (h*w, 3)

    # 聚类为两类（例如：完好 / 损毁）
    kmeans = KMeans(n_clusters=2, random_state=42, n_init='auto')
    labels = kmeans.fit_predict(features)

    # 获取每类像素数，并返回较小类的面积作为损毁面积
    _, counts = np.unique(labels, return_counts=True)
    damage_area = counts.max()
    return int(damage_area)


if __name__ == "__main__":
    image_path = r'C:\Users\86158\Desktop\111111.png'
    image_array = cv2.imread(image_path)

    # 检查是否加载成功
    if image_array is None:
        raise FileNotFoundError(f"图像文件未找到: {image_path}")

    area, contours = cluster_largest_region(image_array)
    print(contours)

    image_draw = image_array.copy()
    cv2.drawContours(image_draw, contours, -1, (0, 255, 0), 2)  # 用绿色画出边界

    output_path = r'C:\Users\86158\Desktop\output_with_contour.png'
    cv2.imwrite(output_path, image_draw)
